import sys, os, pygame
from pygame.locals import *
from pygame.constants import *
from OpenGL.GL import *
from OpenGL.GLU import *
from numpy import *
from PIL import Image
from objloader import *

def getPerspctiveMatrix(fx, fy, cx, cy, w, h, zn, zf):
    return mat([\
    [2*fx/w, 0, 1-2*cx/w, 0],\
    [0, 2*fy/h, 2*cy/h-1, 0],\
    [0, 0, (zf+zn)/(zn-zf), (2*zf*zn)/(zn-zf)],\
    [0, 0, -1, 0]])

def FindExt(path, ext):
	files = os.listdir(path)
	ymls = []
	for file in files:
		if os.path.splitext(file)[1] == ext:
			ymls.append(os.path.join(path, file))
	return ymls

def ReadPose(pose_file):
	poses1 = []
	poses2 = []

	with open(pose_file, 'r') as f:
			content = f.readlines()

	gap = int(content[1])
	nof = (len(content)-2)/gap

	for i in range(0, nof):
		elms = content[i*gap+2].split('\t')
		pose = mat([\
	    [float(elms[1]), float(elms[2]), float(elms[3]), float(elms[10])],\
	    [float(elms[4]), float(elms[5]), float(elms[6]), float(elms[11])],\
	    [float(elms[7]), float(elms[8]), float(elms[9]), float(elms[12])],\
	    [0, 0, 0, 1]])
		poses1.append(pose)

		elms = content[i*gap+2+1].split('\t')
		pose = mat([\
	    [float(elms[1]), float(elms[2]), float(elms[3]), float(elms[10])],\
	    [float(elms[4]), float(elms[5]), float(elms[6]), float(elms[11])],\
	    [float(elms[7]), float(elms[8]), float(elms[9]), float(elms[12])],\
	    [0, 0, 0, 1]])
		poses2.append(pose)
	return [poses1, poses2]

def ReadPose1(pose_file):
	poses = []
	with open(pose_file, 'r') as f:
			content = f.readlines()

	gap = int(content[1])
	nof = (len(content)-2)/gap

	for i in range(0, nof):
		elms = content[i*gap+2].split('\t')
		pose = mat([\
	    [float(elms[1]), float(elms[2]), float(elms[3]), float(elms[10])],\
	    [float(elms[4]), float(elms[5]), float(elms[6]), float(elms[11])],\
	    [float(elms[7]), float(elms[8]), float(elms[9]), float(elms[12])],\
	    [0, 0, 0, 1]])
		poses.append(pose)
	return poses

def ReadPose2(pose_file):
	poses = []
	with open(pose_file, 'r') as f:
			content = f.readlines()

	gap = int(content[1])
	nof = (len(content)-2)/gap

	for i in range(0, nof):
		elms = content[i*gap+2+1].split('\t')
		pose = mat([\
	    [float(elms[1]), float(elms[2]), float(elms[3]), float(elms[10])],\
	    [float(elms[4]), float(elms[5]), float(elms[6]), float(elms[11])],\
	    [float(elms[7]), float(elms[8]), float(elms[9]), float(elms[12])],\
	    [0, 0, 0, 1]])
		poses.append(pose)
	return poses

def RenderImages(obj, pmat1, pmat2, bgimg, dir):
	img_flip = bgimg.transpose(Image.FLIP_TOP_BOTTOM)
	pBits = img_flip.tobytes("raw", "RGB")

	lookAtMatrix = mat([\
	    [1, 0, 0, 0],\
	    [0, -1, 0, 0],\
	    [0, 0, -1, 0],\
	    [0, 0, 0, 1]])


	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
	glDisable(GL_DEPTH_TEST)
	glDrawPixels(bgimg.width, bgimg.height, GL_RGB, GL_UNSIGNED_BYTE, pBits)
	glEnable(GL_DEPTH_TEST)
	 
	glMatrixMode(GL_MODELVIEW)
	mvm = lookAtMatrix*pmat1
	glLoadMatrixf(mvm.T)
	glCallList(obj.gl_list)

	glMatrixMode(GL_MODELVIEW)
	mvm = lookAtMatrix*pmat2
	glLoadMatrixf(mvm.T)
	glCallList(obj.gl_list)

	pygame.image.save(srf, os.path.join(dir, "ar_{0:04d}.png".format(i)))
	glFinish()
	pygame.display.flip()
# python script\pyrender\viewer.py script\pyrender\ape.obj data\rbot\gt_pose.txt data\rbot\frame\ape_regular data\out\tmp

obj_path = sys.argv[1]
pose_path = sys.argv[2]
img_path = sys.argv[3]
out_path = sys.argv[4]

width = 640
height = 512
fx = 650.048
fy = 647.183
cx = 324.328
cy = 257.323
zf = 4000.0
zn = 1.0

pygame.init()
viewport = (width, height)
srf = pygame.display.set_mode(viewport, OPENGL | DOUBLEBUF)

glLightfv(GL_LIGHT0, GL_POSITION,  (-40, 200, 100, 0.0))
#glLightfv(GL_LIGHT0, GL_POSITION,  (0, 100, 0, 0.0))
glLightfv(GL_LIGHT0, GL_AMBIENT, (0.6, 0.6, 0.6, 1.0))
glLightfv(GL_LIGHT0, GL_DIFFUSE, (0.2, 0.2, 0.2, 1.0))
glEnable(GL_LIGHT0)
glEnable(GL_LIGHTING)
glEnable(GL_COLOR_MATERIAL)
glEnable(GL_DEPTH_TEST)
glShadeModel(GL_SMOOTH)

obj = OBJ(obj_path)

kmat = getPerspctiveMatrix(fx, fy, cx, cy, width, height, zn, zf)
glMatrixMode(GL_PROJECTION)
glLoadMatrixf(kmat.T)

imgs = FindExt(img_path, '.png')
poses = ReadPose(pose_path)

clock = pygame.time.Clock()

for i in range(len(poses[0])):
	for e in pygame.event.get():
	    if e.type == QUIT:
	        sys.exit()
	    elif e.type == KEYDOWN and e.key == K_ESCAPE:
	        sys.exit()

	clock.tick(200)
	pose0 = poses[0][i]
	pose1 = poses[1][i]
	pimg = Image.open(imgs[i])

	RenderImages(obj, pose0, pose1, pimg, out_path)

'''
while 1:
	for e in pygame.event.get():
	    if e.type == QUIT:
	        sys.exit()
	    elif e.type == KEYDOWN and e.key == K_ESCAPE:
	        sys.exit()

	clock.tick(30)
	pose0 = poses[0][0]
	pose1 = poses[1][0]
	pimg = Image.open(imgs[0])

	RenderImages(obj, pose0, pose1, pimg, out_path)
'''